import time
from machine import Pin, SPI
from math import ceil, sqrt
from fonts import asc2_0806
import sys
import uos

class TimeoutError(Exception):
    def __init__(self, msg):
        super().__init__(msg)

class Color:
    BLACK = 0x00
    WHITE = 0xff

class Rotate:
    ROTATE_0 = 0
    ROTATE_90 = 1
    ROTATE_180 = 2
    ROTATE_270 = 3

class Screen:
    def __init__(self, width=122,height=250):
        # width=128,height=296
        # width=122, height=250
        self.width = width
        self.height = height
        self.width_bytes = ceil(width / 8)
        self.height_bytes = height

    def __repr__(self):
        print("screen width: %d" % self.screen.width)
        print("screen height: %d" % self.screen.height)
        print("screen width bytes: %d" % self.screen.width_bytes)
        print("screen height bytes: %d" % self.screen.height_bytes)

class Paint:
    def __init__(self, screen=Screen(), rotate=Rotate.ROTATE_90, bg_color=Color.WHITE):
        self.screen = screen
        self.img = bytearray(self.screen.width_bytes * self.screen.height_bytes)
        self.rotate = rotate
        self.bg_color = bg_color

        if self.rotate == Rotate.ROTATE_0 or self.rotate == Rotate.ROTATE_180:
            self.width = self.screen.width
            self.height = self.screen.height
        else:
            self.width = self.screen.height
            self.height = self.screen.width

    def __repr__(self):
        self.screen.__repr__()
        print("rotate: %d" % self.rotate)
        print("background color: 0x%x" % self.bg_color)

    def clear(self, color):
        self.bg_color = color
        for y in range(self.screen.height_bytes):
            for x in range(self.screen.width_bytes):
                addr = x + y * self.screen.width_bytes
                self.img[addr] = self.bg_color

    def _convert_coor(self, x_pos, y_pos, start_from_one=True):
        if self.rotate == Rotate.ROTATE_0:
            x = x_pos
            y = y_pos
        elif self.rotate == Rotate.ROTATE_90:
            x = self.screen.width - y_pos - 1
            y = x_pos
        elif self.rotate == Rotate.ROTATE_180:
            x = self.screen.width - x_pos - 1
            y = self.screen.height - y_pos - 1
        else:
            x = y_pos
            y = self.screen.height - x_pos - 1

        if start_from_one:
            x = x - 1
            y = y - 1

        return x, y

    def draw_point(self, x_pos, y_pos, start_from_one=True):
        x, y = self._convert_coor(x_pos, y_pos)
        if x >= self.screen.width or y >= self.screen.height or x < 0 or y < 0:
            return

        addr = x // 8 + y * self.screen.width_bytes
        if addr >= len(self.img):
            return  # 防止超出范围

        raw = self.img[addr]

        if self.bg_color == Color.WHITE:
            self.img[addr] = raw & ~(0x80 >> (x % 8))
        else:
            self.img[addr] = raw | (0x80 >> (x % 8))

    def draw_line(self, x_start, y_start, x_end, y_end):
        dx = x_end - x_start
        dy = y_end - y_start
        points = []
        if abs(dx) > abs(dy):
            x_inc = (dx > 0) - (dx < 0)
            x_offset = 0
            for _ in range(abs(dx) + 1):
                x_tmp = x_start + x_offset
                y_tmp = y_start + round(dy / dx * x_offset)
                points.append((x_tmp, y_tmp))
                x_offset = x_offset + x_inc
        else:
            y_inc = (dy > 0) - (dy < 0)
            y_offset = 0
            for _ in range(abs(dy) + 1):
                y_tmp = y_start + y_offset
                x_tmp = x_start + round(dx / dy * y_offset)
                points.append((x_tmp, y_tmp))
                y_offset = y_offset + y_inc

        for point in points:
            self.draw_point(point[0], point[1])

    def draw_rectangle(self, x_start, y_start, x_end, y_end):
        self.draw_line(x_start, y_start, x_start, y_end)
        self.draw_line(x_start, y_start, x_end, y_start)
        self.draw_line(x_start, y_end, x_end, y_end)
        self.draw_line(x_end, y_start, x_end, y_end)

    def draw_circle(self, x_center, y_center, radius):
        points = []
        for x in range(x_center - radius, x_center + radius):
            y = y_center + round(sqrt(radius ** 2 - (x - x_center) ** 2))
            points.append((x, y))
            y = y_center - round(sqrt(radius ** 2 - (x - x_center) ** 2))
            points.append((x, y))
        for y in range(y_center - radius, y_center + radius):
            x = x_center + round(sqrt(radius ** 2 - (y - y_center) ** 2))
            points.append((x, y))
            x = x_center - round(sqrt(radius ** 2 - (y - y_center) ** 2))
            points.append((x, y))

        for point in points:
            self.draw_point(point[0], point[1])

    def show_char(self, char, x_start, y_start, font=asc2_0806, font_size=(6, 8), multiplier=1):
        tmp = 0x00
        char_idx = ord(char) - 32
        if multiplier == 1:
            for x_offset in range(font_size[0]):
                tmp = font[char_idx][x_offset]
                for y_offset in range(font_size[1]):
                    if tmp & 0x01:
                        self.draw_point(x_start + x_offset, y_start + y_offset)
                    tmp = tmp >> 1
        else:
            for x_offset in range(font_size[0] * multiplier):
                tmp = font[char_idx][x_offset // multiplier]
                for y_offset in range(font_size[1] * multiplier):
                    if tmp & 0x01:
                        self.draw_point(x_start + x_offset, y_start + y_offset)
                    if y_offset % multiplier == multiplier - 1:
                        tmp = tmp >> 1

    def show_string(self, string, x_start, y_start, font=asc2_0806, font_size=(6, 8), multiplier=1):
        for idx, char in enumerate(string):
            self.show_char(char, x_start + idx * font_size[0] * multiplier, y_start, font, font_size, multiplier)

    def show_string(self, string, x_start, y_start, font=asc2_0806, font_size=(6, 8), multiplier=1):
        for idx, char in enumerate(string):
            self.show_char(char, x_start + idx * font_size[0] * multiplier, y_start, font, font_size, multiplier)

    def show_bitmap(self, bitmap, x_start, y_start, multiplier=1):
        if multiplier == 1:
            for r, row in enumerate(bitmap):
                for c, col in enumerate(row):
                    if col == 1:
                        self.draw_point(x_start + c, y_start + r)
        else:
            for r in range(len(bitmap) * multiplier):
                for c in range(len(bitmap[0] * multiplier)):
                    if bitmap[r // multiplier][c // multiplier] == 1:
                        self.draw_point(x_start + c, y_start + r)

    def draw_bitmap(self, bitmap, x_start, y_start, multiplier=1):
        if isinstance(bitmap, bytearray):
            # Convert bytearray to a list of lists
            width = int(sqrt(len(bitmap) * 8))  # Assuming square bitmap for simplicity
            bitmap = [
                [(bitmap[i // 8] >> (7 - (i % 8))) & 1 for i in range(j * width, (j + 1) * width)]
                for j in range(width)
            ]
        elif not isinstance(bitmap, list):
            raise TypeError("bitmap should be a list or bytearray")

        if multiplier == 1:
            for r, row in enumerate(bitmap):
                for c, col in enumerate(row):
                    if col == 1:
                        self.draw_point(x_start + c, y_start + r)
        else:
            for r in range(len(bitmap) * multiplier):
                for c in range(len(bitmap[0]) * multiplier):
                    if bitmap[r // multiplier][c // multiplier] == 1:
                        self.draw_point(x_start + c, y_start + r)

    def show_img(self, img_path, x_start, y_start):
        raise NotImplementedError

class SSD1680:
    def __init__(self, spi, dc, busy, cs, res,rotate=Rotate.ROTATE_0):
        super().__init__()
        self.spi = spi
        self.dc = dc
        self.busy = busy
        self.cs = cs
        self.res = res
        self.screen = Screen()
        self.paint = Paint(self.screen, rotate, bg_color=Color.WHITE)

        self.dc(1)
        self.chip_sel()

        # Add width and height attributes ufont 使用
        self.width = self.screen.width
        self.height = self.screen.height
        # Initialize buffer ufont 使用
        # self.buffer = self.paint.img

        # 5次局部刷新一次全局更新
        self.partial_update_count = 0
        self.partial_update_threshold = 5  # 根据需要设置阈值

    def chip_sel(self):
        self.cs(0)

    def chip_desel(self):
        self.cs(1)

    def read_busy(self, info="wait busy timeout!", timeout=5):
        st = time.time()
        while self.busy.value() == 1:
            if (time.time() - st) > timeout:
                raise TimeoutError(info)

    def hw_rst(self):
        print("hardware resetting...")
        self.res(0)
        time.sleep(0.2)
        self.res(1)
        time.sleep(0.2)
        self.read_busy("hardware reset timeout!")
        print("hardware reset successful")

    def sw_rst(self):
        print("software resetting...")
        self.write_cmd(0x12)
        self.read_busy("software reset timeout!")
        print("software reset successful")

    def write_cmd(self, cmd: bytearray):
        self.dc(0)
        self.spi.write(cmd.to_bytes(1, 'big'))
        self.dc(1)

    def write_data(self, data: bytearray):
        self.spi.write(data.to_bytes(1, 'big'))

    def init(self):
        self.hw_rst()
        self.sw_rst()

        # deriver output control
        self.write_cmd(0x01)
        self.write_data(0x27)
        self.write_data(0x01)
        self.write_data(0x01)

        # data entry mode
        self.write_cmd(0x11)
        self.write_data(0x01)

        # set ram-x addr start/end pos
        self.write_cmd(0x44)
        self.write_data(0x00)
        self.write_data(0x0F)

        # set ram-y addr start/end pos
        self.write_cmd(0x45)
        self.write_data(0x27)
        self.write_data(0x01)
        self.write_data(0x00)
        self.write_data(0x00)

        # border waveform
        self.write_cmd(0x3c)
        self.write_data(0x05)

        # display update control
        self.write_cmd(0x21)
        self.write_data(0x00)
        self.write_data(0x80)

        # set ram-x addr cnt to 0
        self.write_cmd(0x4e)
        self.write_data(0x00)
        # set ram-y addr cnt to 0x127
        self.write_cmd(0x4F)
        self.write_data(0x27)
        self.write_data(0x01)

    # reset ram 的写入xy初始坐标
    def resetXY(self):
        # set ram-x addr cnt to 0
        self.write_cmd(0x4e)
        self.write_data(0x00)
        # set ram-y addr cnt to 0x127
        self.write_cmd(0x4F)
        self.write_data(0x27)
        self.write_data(0x01)

    # 加载图像到ram
    def update_mem(self):
        print("updating the memory...")
        self.write_cmd(0x24)

        # Clear the buffer before updating
        # self.buffer = bytearray(self.width * self.height // 8)

        for k in range(self.paint.screen.height_bytes * self.paint.screen.width_bytes):
            byte = self.paint.img[k]
            self.write_data(byte)
        print("updating memory successful")

    # ram 图像显示到屏幕
    def update_screen(self):
        # display update control
        self.write_cmd(0x22)
        # /* 显示模式设置指令，参考SSD1680 datasheet */
        # /* 在显示之前，初始化时钟，DC-DC，加载LUT，进入全刷模式，执行图像刷新，关闭DC-DC时钟 */
        # 局部刷新
        # self.write_data(0xFF) 
        # 全局刷新
        self.write_data(0xF7) 
        # update
        print("updating the screen...")
        self.write_cmd(0x20)
        # /* 执行指令，参考SSD1680 datasheet */
        self.read_busy("update screen timeout!")
        print("update screen successful")

    # Clear register 0x26
    def clear0x26(self):
        self.write_cmd(0x26)  # Write to RAM command for register 0x26
        for _ in range(self.paint.screen.width_bytes * self.paint.screen.height):
            self.write_data(0xFF)  # Write white (0xFF) to all pixels
        self.read_busy("Clear register 0x26 timeout!")
        print("Register 0x26 cleared successfully")

    # 局部清除img
    def clear_partial(self, x_start, y_start, x_end, y_end):
        for y in range(y_start, y_end):
            for x in range(x_start, x_end):
                addr = x // 8 + y * self.paint.screen.width_bytes
                if addr < len(self.paint.img):
                    self.paint.img[addr] = Color.WHITE

    # 局部更新
    def update_partial(self, x_start, y_start, x_end, y_end,update_mode=0XFF):
        # self.clear_partial(x_start, y_start, x_end, y_end)
        self.clear0x26()

        # 重置xy
        # 重置 RAM X 和 Y 地址计数器
        self.resetXY()
        # 加载整个图像到 RAM
        self.write_cmd(0x24)  # 写入 RAM 命令
        for byte in self.paint.img:
            self.write_data(byte)

        # # 重置 RAM X 和 Y 地址计数器到 x_start 和 y_start
        # self.write_cmd(0x4e)  # 设置 RAM X 地址计数器
        # self.write_data(x_start // 8)  # 按字节设置 X 地址计数器
        # self.write_cmd(0x4F)  # 设置 RAM Y 地址计数器
        # self.write_data(y_start & 0xFF)
        # self.write_data((y_start >> 8) & 0xFF)

        # # 加载指定范围的图像到 RAM
        # self.write_cmd(0x24)  # 写入 RAM 命令
        # for y in range(y_start, y_end):
        #     for x in range(x_start, x_end):
        #         addr = x // 8 + y * self.paint.screen.width_bytes
        #         bit_offset = x % 8
        #         if addr < len(self.paint.img):
        #             byte = self.paint.img[addr]
        #             # 处理字节内的偏移
        #             mask = 0x80 >> bit_offset
        #             byte = (byte & mask) | (self.paint.img[addr] & ~mask)
        #             self.write_data(byte)

        # 更新部分显示
        self.write_cmd(0x22)  # 图像更新
        self.write_data(update_mode)  # 局部更新 91 99 B1 B9 C7 CF F7 FF
        self.write_cmd(0x20)  # 等待忙信号变为低电平
        self.read_busy("partial update timeout!", timeout=10)  # 将超时时间增加到 10 秒
        print("partial update successful")
        self.resetXY()
        
    
    def update(self):
        self.update_mem()
        # self.clear0x26()
        self.update_screen()
        self.resetXY()

    def clear(self, *args, **kwargs):
        self.paint.clear(*args, **kwargs)

    def draw_point(self, *args, **kwargs):
        self.paint.draw_point(*args, **kwargs)

    def draw_line(self, *args, **kwargs):
        self.paint.draw_line(*args, **kwargs)

    def draw_rectangle(self, *args, **kwargs):
        self.paint.draw_rectangle(*args, **kwargs)

    def draw_circle(self, *args, **kwargs):
        self.paint.draw_circle(*args, **kwargs)

    def show_char(self, *args, **kwargs):
        self.paint.show_char(*args, **kwargs)

    def show_string(self, *args, **kwargs):
        self.paint.show_string(*args, **kwargs)

    def show_bitmap(self, *args, **kwargs):
        self.paint.show_bitmap(*args, **kwargs)

    def draw_bitmap(self, *args, **kwargs):
        self.paint.draw_bitmap(*args, **kwargs)

    def show_img(self, *args, **kwargs):
        self.paint.show_img(*args, **kwargs)

    # 测试局部更新
    def test_partial_update(self,update_mode):
        # Clear the entire screen
        self.clear(Color.WHITE)
        self.update()
        self.resetXY()

   
        # Write some text in the middle
        self.show_string("11:11", 0, 20,multiplier=2)
        self.update()

        # Initialize time
        hour = 11
        minute = 11

        while True:
            # Increment minute
            minute += 1
            if minute == 60:
                minute = 0
                hour += 1
            if hour == 24:
                hour = 0

            # Format time string
            time_str = f"{hour:02d}:{minute:02d}"


            self.partial_update_count += 1
            if self.partial_update_count >= self.partial_update_threshold:
                self.resetXY()
                self.clear(Color.WHITE)
                self.show_string(time_str, 0, 20, multiplier=2)  # Changed y position from 10 to 20
                self.update()
                self.partial_update_count = 0
            else:
                # Clear previous time
                self.clear_partial(0, 20, 100, 50)  # Adjust coordinates as needed
                self.update_partial(0, 20, 100, 50, update_mode)
                # Write new time
                self.show_string(time_str, 0, 20, multiplier=2)  # Changed y position from 10 to 20
                # Perform partial update
                self.update_partial(0, 20, 100, 50, update_mode)

            # Wait for a minute (or less for demonstration purposes)
            time.sleep(1)  # Adjust as needed

       





